/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2017-2020. All rights reserved.
 * Description: hiai model manager
 */
#include <map>
#include <string>
#include <vector>
#include <cstdint>
#include <cstring>
#include <cstdlib>
#include <chrono>
#include <thread>
#include <iostream>
#include "securec.h"
#include "tensor/hiai_nd_tensor_buffer_impl.h"
#include "tensor/hiai_nd_tensor_desc_impl.h"

using namespace std;

#ifdef __cplusplus
extern "C" {
#endif
struct HIAI_NDTensorDesc {
    int32_t* dims;
    size_t dimNum;
    HIAI_DataType dataType;
    HIAI_Format format;
};

struct HIAI_MR_NDTensorBuffer {
    HIAI_NDTensorDesc* desc;
    size_t size;
    void* data;
    void* handle;
    bool owner;
};

static size_t g_dataTypeElementSize[HIAI_DATATYPE_DOUBLE + 1] = {
    sizeof(uint8_t), sizeof(float), sizeof(float) / 2, sizeof(int32_t), sizeof(int8_t),
    sizeof(int16_t), sizeof(char), sizeof(int64_t), sizeof(uint32_t), sizeof(double)
};

size_t HIAI_NDTensorDesc_GetElementNum_Impl(const HIAI_NDTensorDesc* tensorDesc)
{
    if (tensorDesc == NULL) {
        return 0;
    }

    size_t totalSize = 1;
    for (size_t i = 0; i < tensorDesc->dimNum; i++) {
        int32_t dimValue = tensorDesc->dims[i];
        totalSize *= dimValue;
    }

    return totalSize;
}

size_t HIAI_NDTensorDesc_GetByteSize_Impl(const HIAI_NDTensorDesc* tensorDesc)
{
    size_t totalDimNum = HIAI_NDTensorDesc_GetElementNum_Impl(tensorDesc);
    if (totalDimNum == 0) {
        return 0;
    }

    return totalDimNum * g_dataTypeElementSize[tensorDesc->dataType];
}

HIAI_MR_NDTensorBuffer* HIAI_MR_NDTensorBuffer_CreateFromNDTensorDesc_Impl(const HIAI_NDTensorDesc* desc)
{
    // 复制一份desc
    HIAI_NDTensorDesc* d = HIAI_NDTensorDesc_Clone_Impl(desc);
    if (d == NULL) {
        return NULL;
    }

    // 打桩创建rom侧buffer
    HIAI_MR_NDTensorBuffer* buffer = static_cast<HIAI_MR_NDTensorBuffer*>(malloc(sizeof(HIAI_MR_NDTensorBuffer)));
    (void)memset_s(buffer, sizeof(HIAI_MR_NDTensorBuffer), 0, sizeof(HIAI_MR_NDTensorBuffer));
    int32_t totalSize = HIAI_NDTensorDesc_GetByteSize_Impl(desc);
    buffer->desc = d;
    buffer->size = totalSize;
    buffer->data = malloc(totalSize);
    buffer->handle = nullptr;
    buffer->owner = true;

    // 进行handle封装
    HIAI_MR_NDTensorBuffer* ndBuffer = (HIAI_MR_NDTensorBuffer*)malloc(sizeof(HIAI_MR_NDTensorBuffer));
    (void)memset_s(ndBuffer, sizeof(HIAI_MR_NDTensorBuffer), 0, sizeof(HIAI_MR_NDTensorBuffer));
    ndBuffer->size = buffer->size;
    ndBuffer->data = buffer->data;
    ndBuffer->handle = buffer;
    ndBuffer->owner = true;
    ndBuffer->desc = HIAI_NDTensorDesc_Clone_Impl(desc);

    return ndBuffer;
}

HIAI_NDTensorDesc* HIAI_NDTensorDesc_Create_Impl(
    const int32_t* dims, size_t dimNum, HIAI_DataType dataType, HIAI_Format format)
{
    if (dimNum == 0 || dimNum > (__UINT32_MAX__ / sizeof(int32_t))) {
        return NULL;
    }

    HIAI_NDTensorDesc* desc = static_cast<HIAI_NDTensorDesc*>(malloc(sizeof(HIAI_NDTensorDesc)));
    size_t dimsSize = dimNum * sizeof(int32_t);
    desc->dims = (int32_t*)malloc(dimsSize);
    (void)memcpy_s(desc->dims, dimsSize, dims, dimsSize);
    desc->dimNum = dimNum;
    desc->dataType = dataType;
    desc->format = format;
    return desc;
}

HIAI_NDTensorDesc* HIAI_NDTensorDesc_Clone_Impl(const HIAI_NDTensorDesc* tensorDesc)
{
    if (tensorDesc == NULL) {
        return NULL;
    }

    return HIAI_NDTensorDesc_Create_Impl(
        tensorDesc->dims, tensorDesc->dimNum, tensorDesc->dataType, tensorDesc->format);
}

HIAI_MR_NDTensorBuffer* HIAI_MR_NDTensorBuffer_CreateFromSize_Impl(const HIAI_NDTensorDesc* desc, size_t size)
{
    // 复制一份desc
    HIAI_NDTensorDesc* d = HIAI_NDTensorDesc_Clone_Impl(desc);

    // 打桩创建rom侧buffer
    HIAI_MR_NDTensorBuffer* buffer = static_cast<HIAI_MR_NDTensorBuffer*>(malloc(sizeof(HIAI_MR_NDTensorBuffer)));
    (void)memset_s(buffer, sizeof(HIAI_MR_NDTensorBuffer), 0, sizeof(HIAI_MR_NDTensorBuffer));
    buffer->desc = d;
    buffer->size = size;
    buffer->data = malloc(size);
    buffer->handle = nullptr;
    buffer->owner = true;

    // 进行handle封装
    HIAI_MR_NDTensorBuffer* ndBuffer = (HIAI_MR_NDTensorBuffer*)malloc(sizeof(HIAI_MR_NDTensorBuffer));
    (void)memset_s(ndBuffer, sizeof(HIAI_MR_NDTensorBuffer), 0, sizeof(HIAI_MR_NDTensorBuffer));
    ndBuffer->size = buffer->size;
    ndBuffer->data = buffer->data;
    ndBuffer->handle = buffer;
    ndBuffer->owner = true;
    ndBuffer->desc = HIAI_NDTensorDesc_Clone_Impl(desc);

    return ndBuffer;
}

void* HIAI_MR_NDTensorBuffer_GetHandle_Impl(const HIAI_MR_NDTensorBuffer* ndBuffer)
{
    if (ndBuffer == NULL) {
        return NULL;
    }
    return ndBuffer->handle;
}

size_t HIAI_MR_NDTensorBuffer_GetSize_Impl(const HIAI_MR_NDTensorBuffer* ndBuffer)
{
    return 4;
}

HIAI_NDTensorDesc* HIAI_MR_NDTensorBuffer_GetNDTensorDesc_Impl(const HIAI_MR_NDTensorBuffer* tensorBuffer)
{
    if (tensorBuffer == NULL) {
        return 0;
    }
    return tensorBuffer->desc;
}

HIAI_MR_NDTensorBuffer* HIAI_MR_NDTensorBufferNoCopy(const HIAI_NDTensorDesc* desc,
    size_t size, const void* data)
{
    HIAI_MR_NDTensorBuffer* ndBuffer = (HIAI_MR_NDTensorBuffer*)malloc(sizeof(HIAI_MR_NDTensorBuffer));
    if (ndBuffer == NULL) {
        return NULL;
    }
    (void)memset_s(ndBuffer, sizeof(HIAI_MR_NDTensorBuffer), 0, sizeof(HIAI_MR_NDTensorBuffer));

    ndBuffer->desc = HIAI_NDTensorDesc_Clone_Impl(desc);
    if (ndBuffer->desc == NULL) {
        free(ndBuffer);
        return NULL;
    }

    ndBuffer->size = size;
    ndBuffer->data = (void*)data;
    ndBuffer->handle = NULL;
    ndBuffer->owner = false;
    return ndBuffer;
}

HIAI_MR_NDTensorBuffer* HIAI_MR_NDTensorBuffer_CreateNoCopy_Impl(
    const HIAI_NDTensorDesc* desc, const void* data, size_t dataSize)
{
    if (desc == NULL) {
        return NULL;
    }

    if (data == NULL) {
        return NULL;
    }

    if (dataSize == 0) {
        return NULL;
    }

    return HIAI_MR_NDTensorBufferNoCopy(desc, dataSize, data);
}

void HIAI_NDTensorDesc_Destroy_Impl(HIAI_NDTensorDesc** tensorDesc)
{
    if (tensorDesc == NULL || *tensorDesc == NULL) {
        return;
    }
    if ((*tensorDesc)->dims != NULL) {
        free((*tensorDesc)->dims);
    }
    free((*tensorDesc));
    (*tensorDesc) = NULL;
}

void* HIAI_MR_NDTensorBuffer_GetData_Impl(const HIAI_MR_NDTensorBuffer* ndBuffer)
{
    if (ndBuffer == NULL) {
        return NULL;
    }
    return ndBuffer->data;
}

void HIAI_MR_NDTensorBuffer_ReleaseHandle(HIAI_MR_NDTensorBuffer** buffer)
{
    if (buffer == NULL || *buffer == NULL) {
        return;
    }
    // 释放buffer desc
    HIAI_NDTensorDesc* desc = (*buffer)->desc;
    HIAI_NDTensorDesc_Destroy_Impl(&desc);
    // 释放buffer data
    free((*buffer)->data);
    // 释放buffer
    free(*buffer);
    *buffer = NULL;
}

void HIAI_MR_NDTensorBuffer_Destroy_Impl(HIAI_MR_NDTensorBuffer** ndBuffer)
{
    if (ndBuffer == NULL || *ndBuffer == NULL) {
        return;
    }
    // 释放ndbuffer desc
    HIAI_NDTensorDesc* desc = (*ndBuffer)->desc;
    HIAI_NDTensorDesc_Destroy_Impl(&desc);

    // 释放ndbuffer handle
    HIAI_MR_NDTensorBuffer* buffer = (HIAI_MR_NDTensorBuffer*)(HIAI_MR_NDTensorBuffer_GetHandle_Impl(*ndBuffer));
    HIAI_MR_NDTensorBuffer_ReleaseHandle(&buffer);

    // 释放ndbuffer
    free(*ndBuffer);
    *ndBuffer = NULL;
}

HIAI_MR_NDTensorBuffer* HIAI_MR_NDTensorBuffer_CreateFromBuffer_Impl(
    const HIAI_NDTensorDesc* desc, const void* data, size_t dataSize)
{
    HIAI_MR_NDTensorBuffer* buffer = HIAI_MR_NDTensorBuffer_CreateFromNDTensorDesc_Impl(desc);
    if (buffer == NULL) {
        return NULL;
    }

    if (buffer->size != dataSize) {
        HIAI_MR_NDTensorBuffer_Destroy_Impl(&buffer);
        return NULL;
    }

    if (memcpy_s(buffer->data, buffer->size, data, dataSize) != 0) {
        HIAI_MR_NDTensorBuffer_Destroy_Impl(&buffer);
        return NULL;
    }
    return buffer;
}

int32_t HIAI_MR_NDTensorBuffer_GetFd_Impl(const HIAI_MR_NDTensorBuffer* tensorBuffer)
{
    std::cout << __func__ << std::endl;

    return 0;
}

int32_t HIAI_MR_NDTensorBuffer_GetOriginFd_Impl(const HIAI_MR_NDTensorBuffer* tensorBuffer)
{
    std::cout << __func__ << std::endl;

    return -1;
}

const int32_t* HIAI_NDTensorDesc_GetDims_Impl(const HIAI_NDTensorDesc* tensorDesc)
{
    if (tensorDesc == nullptr) {
        return nullptr;
    }
    return tensorDesc->dims;
}

#ifdef __cplusplus
}
#endif
